/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2024-2024.
 * Description: support rtos kexec enhance feature
 * Author: wanglinhui
 * Create: 2024-02-17
 */
#include <linux/mm.h>
#include <linux/slab.h>
#include <linux/init.h>
#include <linux/kexec.h>
#include <linux/ima.h>
#include <linux/kernel.h>
#include <asm/setup.h>

#define KEXEC_BOOT_FLAG "kexec_boot"

static bool match_param(const char *input_param)
{
	static const char * const param[] = { "initrd", "root" };
	int i;

	for (i = 0; i < ARRAY_SIZE(param); i++) {
		if (strcmp(input_param, param[i]) == 0)
			return true;
	}

	return false;
}

static int remove_cmdline_param(char *dst_cmdline_buf, size_t dst_buf_len,
	const char *src_cmdline)
{
	int k = 0;
	char *param;
	char *val;
	char *args;
	char *args_begin;
	char *tmp_cmdline;

	if (dst_buf_len < strlen(src_cmdline) + 1)
		return -EINVAL;

	tmp_cmdline = kstrdup(src_cmdline, GFP_KERNEL);
	if (!tmp_cmdline)
		return -ENOMEM;

	args = skip_spaces(tmp_cmdline);
	while (*args) {
		args_begin = args;
		args = next_arg(args, &param, &val);

		if (match_param(param))
			continue;

		memcpy(dst_cmdline_buf + k, src_cmdline + (args_begin - tmp_cmdline),
			args - args_begin);
		k = k + args - args_begin;
	}

	dst_cmdline_buf[k] = '\0';
	kfree(tmp_cmdline);
	return 0;
}

static int append_kexec_boot_flag(struct kimage *image, int kernel_fd)
{
	char *ptr = image->cmdline_buf;

	/*
	 * Concatenating two strings requires a space character,
	 * and then a '\0' needs to be appended to the end of the string,
	 * thus increasing the length of 2 characters.
	 */
	image->cmdline_buf_len = strlen(ptr) + strlen(KEXEC_BOOT_FLAG) + 2;
	if (image->cmdline_buf_len > COMMAND_LINE_SIZE) {
		pr_err("cmdline has no space to append %s flag", KEXEC_BOOT_FLAG);
		image->cmdline_buf_len = 0;
		return -EINVAL;
	}

	image->cmdline_buf = kmalloc(image->cmdline_buf_len, GFP_KERNEL);
	if (!image->cmdline_buf) {
		image->cmdline_buf_len = 0;
		return -ENOMEM;
	}

	sprintf(image->cmdline_buf, "%s %s", ptr, KEXEC_BOOT_FLAG);
	ima_kexec_cmdline(kernel_fd, image->cmdline_buf,
				image->cmdline_buf_len - 1);

	kfree(ptr);
	return 0;
}

int kexec_cmline_process(struct kimage *image, int kernel_fd)
{
	char *dst_buf;
	size_t dst_buf_len;
	int ret;

	dst_buf_len = strlen(saved_command_line) + 1;
	dst_buf = kzalloc(dst_buf_len, GFP_KERNEL);
	if (!dst_buf) {
		image->cmdline_buf_len = 0;
		return -ENOMEM;
	}

	ret = remove_cmdline_param(dst_buf, dst_buf_len, saved_command_line);
	if (ret != 0) {
		image->cmdline_buf_len = 0;
		kfree(dst_buf);
		return ret;
	}

	image->cmdline_buf = kstrdup(dst_buf, GFP_KERNEL);
	kfree(dst_buf);
	if (!image->cmdline_buf) {
		image->cmdline_buf_len = 0;
		return -ENOMEM;
	}

	if (strstr(image->cmdline_buf, KEXEC_BOOT_FLAG)) {
		image->cmdline_buf_len = strlen(image->cmdline_buf) + 1;
		return 0;
	}

	ret = append_kexec_boot_flag(image, kernel_fd);
	if (ret != 0) {
		kfree(image->cmdline_buf);
		image->cmdline_buf = NULL;
		return ret;
	}

	return 0;
}

static int __init kexec_boot_setup(char *str)
{
	pr_info("[kexec_hardening] current kernel boot from kexec");
	return 1;
}
__setup(KEXEC_BOOT_FLAG, kexec_boot_setup);
